package com.congee02.graph;

import java.util.*;

public class CountPairs {

    public long countPairs(int n, int[][] edges) {

        if (edges.length == 0) {
            return (long) n * (long) (n - 1) / 2;
        }

        int[] parent = new int[n];
        Arrays.setAll(parent, i -> i);

        int[] size = new int[n];
        Arrays.fill(size, 1);

        for (int[] edge : edges) {
            union(size, parent, edge[0], edge[1]);
        }

        long ans = 0;
        long sum = 0;
        for (int i = 0; i < n; i ++ ) {
            if (parent[i] == i) {
                ans += size[i] * sum;
                sum += size[i];
            }
        }

        return ans;
    }

    private int root(int[] parent, int x) {
        return parent[x] == x ? x : (parent[x] = root(parent, parent[x]));
    }

    private void union(int[] size, int[] parent, int x, int y) {

        int rx = root(parent, x);
        int ry = root(parent, y);

        if (rx == ry) {
            return;
        }

        // merge by rank
        // if (size[rx] > size[ry]) {
        //     int t = rx;
        //     rx = ry;
        //     ry = t;
        // }

        parent[rx] = ry;
        size[ry] += size[rx];
    }

}
